from functools import wraps
from flask import session, redirect, url_for, flash
from models.user import User, db

def login_required(f):
    """
    装饰器，用于保护需要登录才能访问的路由
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            flash('请先登录')
            return redirect(url_for('auth.login'))
        return f(*args, **kwargs)
    return decorated_function

def admin_required(f):
    """
    装饰器，用于保护需要管理员权限才能访问的路由
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # 检查用户是否已登录
        if 'user_id' not in session:
            flash('请先登录')
            return redirect(url_for('auth.login'))
        
        # 获取当前用户
        user_id = session.get('user_id')
        user = User.query.get(user_id)
        
        # 检查用户是否存在
        if not user:
            flash('用户不存在')
            return redirect(url_for('auth.login'))
        
        # 在这个简单的示例中，我们假设用户ID为1的是管理员
        # 或者用户名为'admin'的用户也是管理员
        if user_id != 1 and user.username != 'admin':
            flash('需要管理员权限')
            return redirect(url_for('dashboard'))
        
        return f(*args, **kwargs)
    return decorated_function